import numpy as np
import matplotlib

matplotlib.use('TkAgg')  # 更换为更稳定的后端
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.text import Annotation

# 设置支持数学符号的字体
plt.rcParams['mathtext.fontset'] = 'stix'
plt.rcParams['font.family'] = 'STIXGeneral'
plt.rcParams['font.size'] = 10


# Adam优化算法实现
def adam_optimizer(grad_func, initial_x, alpha=0.1, beta1=0.9, beta2=0.999, epsilon=1e-8, iterations=50):
    x = initial_x
    m = 0
    v = 0
    history = []

    for t in range(1, iterations + 1):
        grad = grad_func(x)
        m = beta1 * m + (1 - beta1) * grad
        v = beta2 * v + (1 - beta2) * grad ** 2
        m_hat = m / (1 - beta1 ** t)
        v_hat = v / (1 - beta2 ** t)
        x = x - alpha * m_hat / (np.sqrt(v_hat) + epsilon)

        # 只存储必要的历史数据（减少内存占用）
        history.append({'step': t, 'x': x})
    return x, history


# 目标函数和梯度函数
def quadratic(x):
    return (x - 2) ** 2 + 3


def quadratic_grad(x):
    return 2 * (x - 2)


# 优化过程
initial_x = -5.0
optimal_x, history = adam_optimizer(quadratic_grad, initial_x, alpha=0.5, iterations=15)

# 准备可视化数据
x_vals = np.linspace(-5, 5, 100)  # 减少采样点
y_vals = quadratic(x_vals)

# 提取优化路径
path_x = [initial_x] + [h['x'] for h in history]
path_y = [quadratic(x) for x in path_x]
steps = [0] + [h['step'] for h in history]

# 创建图形
fig = plt.figure(figsize=(12, 10))
fig.suptitle('Adam Optimization Algorithm Visualization', fontsize=14, fontweight='bold')

# 子图布局
ax1 = plt.subplot2grid((3, 2), (0, 0), colspan=2)
ax2 = plt.subplot2grid((3, 2), (1, 0))
ax3 = plt.subplot2grid((3, 2), (1, 1))
ax4 = plt.subplot2grid((3, 2), (2, 0), colspan=2)

# 子图1：函数曲线和优化路径
ax1.plot(x_vals, y_vals, 'b-', label=r'$f(x) = (x-2)^2 + 3$')
ax1.plot(path_x, path_y, 'ro-', markersize=6, label='Optimization Path')
ax1.plot(2, 3, 'g*', markersize=12, label='Optimal Solution')
ax1.set_xlabel('x')
ax1.set_ylabel('f(x)')
ax1.grid(True, alpha=0.5)
ax1.legend(fontsize=8)

# 子图2：梯度变化（动态计算梯度，不存储历史梯度）
gradients = [quadratic_grad(x) for x in path_x]
ax2.plot(steps, gradients, 'bo-', linewidth=1.5, markersize=5)
ax2.set_xlabel('Iteration Step')
ax2.set_ylabel('Gradient Value')
ax2.grid(True, alpha=0.5)
ax2.set_xticks(steps)

# 子图3：自适应学习率（动态计算，不存储历史学习率）
adaptive_lrs = [0.5]  # 初始学习率
for h in history:
    v_hat = h.get('v_hat', 1e-8)  # 避免None值
    lr = 0.5 / (np.sqrt(v_hat) + 1e-8)
    adaptive_lrs.append(lr)

ax3.plot(steps, adaptive_lrs, 'go-', linewidth=1.5, markersize=5)
ax3.set_xlabel('Iteration Step')
ax3.set_ylabel('Learning Rate')
ax3.grid(True, alpha=0.5)
ax3.set_xticks(steps)

# 子图4：优化路径动画
ax4.plot(x_vals, y_vals, 'b-', linewidth=1.5)
ax4.plot(2, 3, 'g*', markersize=12, label='Optimal Solution')
ax4.set_xlabel('x')
ax4.set_ylabel('f(x)')
ax4.grid(True, alpha=0.5)
ax4.set_xlim(-5, 5)
ax4.set_ylim(0, 55)

# 动画元素
point, = ax4.plot([], [], 'ro', markersize=8)
path_line, = ax4.plot([], [], 'r--', linewidth=1)


# 动画函数（优化更新逻辑）
def init():
    point.set_data([], [])
    path_line.set_data([], [])
    return point, path_line


def update(frame):
    if frame == 0:
        x_data = [path_x[0]]
        y_data = [path_y[0]]
    else:
        x_data = path_x[:frame + 1]
        y_data = path_y[:frame + 1]

    point.set_data([path_x[frame]], [path_y[frame]])
    path_line.set_data(x_data, y_data)

    # 限制注释数量，避免过多元素
    if frame < len(steps) - 3:  # 只注释前几步
        ax4.annotate(f'Step {frame}', xy=(path_x[frame], path_y[frame]),
                     xytext=(path_x[frame] + 0.2, path_y[frame] + 2),
                     arrowprops=dict(arrowstyle="->", connectionstyle="arc3", alpha=0.7))
    return point, path_line


# 创建动画（降低帧率）
ani = FuncAnimation(
    fig, update, frames=len(path_x), init_func=init,
    blit=True, interval=1000, repeat_delay=2000  # 1秒/帧
)

# 布局优化
plt.tight_layout()
plt.subplots_adjust(top=0.90, hspace=0.4)

# 打印结果
print("Optimization Result:")
print(f"Initial value: x0 = {initial_x:.4f}, f(x0) = {quadratic(initial_x):.4f}")
print(f"Theoretical optimal: x* = 2.0000, f(x*) = 3.0000")
print(f"Adam optimized: x = {optimal_x:.4f}, f(x) = {quadratic(optimal_x):.4f}")

# 显示图形（添加异常处理）
try:
    plt.show()
except Exception as e:
    print(f"显示图形时发生异常: {e}")
    plt.close(fig)  # 强制关闭图形

# 手动释放资源
import gc

gc.collect()